/*
// <Alex Langhart Comment>

    user_level_test_program.c

// </Alex Langhart Comment>
*/

#include <stdio.h>
#include <sys/types.h>
#include <unistd.h>
#include <linux/sched.h>
#include <linux/unistd.h>

#define __NR_sched_setaffinity 241
#define __NR_sched_getaffinity 242
#define __NR_get_wrr_weight 337
#define __NR_set_wrr_scheduler 338

#define WEIGHT_1 3
#define WEIGHT_2 1

pid_t fork_busywork(char** argv);

int main(int argc, char** argv){

    /*
    // <Alex Langhart Comment>
    
    Write a user-level test program. The program first forks multiple processes and
    issues the set_wrr_scheduler() system call to configure these processes to be scheduled
    under the WRR policy with different weights. To verify that these processes are
    running with WRR, the test program issues the get_wrr_weight() system call to query
    individual processes' weights. At last, the program set the CPU affinity of these
    processes (using sched_setaffinity()) to the same core so as to test if the WRR policy
    works.
    
    // </Alex Langhart Comment>
    */
    
    pid_t self_pid = getpid();

    pid_t  pid1 = -1, pid2 = -1;
    int weight1 = 0;
    int weight2 = 0;
    unsigned long affinity1 = 1L;
    unsigned long affinity2 = 1L;

    pid1 = fork_busywork(argv);
    pid2 = fork_busywork(argv);

    printf("Main process: main pid=%d\n", self_pid);
    printf("Main process: processes created:%d and %d\n", pid1, pid2);

    printf("Main process: Switching to WRR and setting weights to %d and %d\n", WEIGHT_1, WEIGHT_2);
    syscall(__NR_set_wrr_scheduler, pid1, WEIGHT_1);
    syscall(__NR_set_wrr_scheduler, pid2, WEIGHT_2);

    printf("Main process: querying the kernel for the process weights...\n");
    syscall(__NR_get_wrr_weight, pid1, &weight1);
    syscall(__NR_get_wrr_weight, pid2, &weight2);

    printf("Main process: Kernel reports weight1=%d, weight2=%d\n", weight1, weight2);

    if(weight1 != WEIGHT_1 || weight2 != WEIGHT_2){
        printf("Main process: weights do not match!!\n");
    }

    printf("Main process: Setting CPU affinities to 1\n");
    syscall(__NR_sched_setaffinity, pid1, sizeof(unsigned long), &affinity1);
    syscall(__NR_sched_setaffinity, pid2, sizeof(unsigned long), &affinity2);

    affinity1 = affinity2 = -1;
    syscall(__NR_sched_getaffinity, pid1, sizeof(unsigned long), &affinity1);
    syscall(__NR_sched_getaffinity, pid2, sizeof(unsigned long), &affinity2);
    printf("Main process: Returned affinities are %d=%lu and %d=%lu\n", pid1, affinity1, pid2, affinity2);

    printf("Main process: Entering infinite sleep loop (use this time to verify CPU utilization)...\n");
    while(1){
        sleep(1000);    
    }

    return 0;
    
}

pid_t fork_busywork(char** argv){
    pid_t pid = fork();

    if(pid == 0){
        execve("./busywork", argv, NULL);
    }
    return pid;
}
